import torch
import torch.nn as nn
from matplotlib import pyplot as plt

from src.encoder import Encoder
from src.decoder import decoder

class DECSNet(nn.Module):
    def __init__(self, in_channel, out_channel, resolution, patchsz, mode, pretrained, cnn_dim, transform_dim, num_class):
        super().__init__()
        self.encoder = Encoder(in_channel, out_channel, resolution, patchsz, mode, pretrained, cnn_dim, transform_dim)
        self.decoder = decoder(num_class, transform_dim)

    def forward(self, x):
        # print(x.shape)
        # test时加上这行代码
        # x = x.repeat(1, 3, 1, 1)
        o1, o2, o3, o4 = self.encoder(x)
        out = self.decoder(o1, o2, o3, o4)
        return out
